import torch
import torch.nn as nn
import torchvision

import pandas as pd
import numpy as np


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=(3, 3)),
            nn.ReLU(),
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=(3, 3)),
            nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.Linear(in_features=16 * 6,out_features=24),
            nn.Linear(in_features=24, out_features=10)
        )

    def forward(self, x):
        x = self.layer1(x)
        out = self.layer2(x)
        return out


if __name__ == '__main__':
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    model = Model().to(device)

    train_set = pd.read_csv('mnist/train-images-idx3-ubyte')
    print(train_set)
